#导入库
import cv2
import dlib
import os
import shutil
import numpy as np
from tkinter import *
from tkinter import messagebox
import tkinter.filedialog as fd

#关键点编码
def encoder_face(image,detector,predictor,encoder,upsample=1,jet=1):
    #检测人脸
    faces = detector(image,upsample)
    #对每张人脸进行关键点检测
    # gray = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)#灰度转换
    faces_keypoints = [ predictor(image,face) for face in faces ]#每张人脸关键点
    return [ np.array(encoder.compute_face_descriptor(image,face_keypoint,jet)) for face_keypoint in faces_keypoints ]

#人脸比较（通过欧式距离）
def compare_faces(face_encoding, test_encoding):
    return list(np.linalg.norm(np.array(face_encoding) - np.array(test_encoding),axis=1))#按行向量处理

#人脸比较，输出对应名称
def compare_faces_order(face_encoding, test_encoding, names):
    distance = list(np.linalg.norm(np.array(face_encoding) - np.array(test_encoding),axis=1))
    return zip(*sorted(zip(distance, names)))

#读取图片
def image_read(path):
    imgs = []
    imagePaths = [os.path.join(path, f) for f in os.listdir(path)]
    for imagePath in imagePaths:
        img = cv2.imread(imagePath)
        img = img[:, :, ::-1]
        # print(imagePath)
        imgs.append(img)
    # cv2.imshow('1',imgs[0])
    # cv2.waitKey(0)
    return imgs

def openFloder1():
    global path_eg
    path_eg= fd.askdirectory()  # 打开文件
    print(path_eg)
    show_folderPath1.delete(0, END)  # 清空
    show_folderPath1.insert(0, path_eg)  # 写入路径

def openFloder2():
    global path_rs
    path_rs= fd.askdirectory()  # 打开文件
    print(path_rs)
    show_folderPath2.delete(0, END)  # 清空
    show_folderPath2.insert(0, path_rs)  # 写入路径

def openFloder3():
    global path_cl
    path_cl= fd.askdirectory()  # 打开文件
    print(path_cl)
    show_folderPath3.delete(0, END)  # 清空
    show_folderPath3.insert(0, path_cl)  # 写入路径

def main():
    #读取分类依据图片
    #path_eg = 'first'#这个部分要让用户进行选择
    #path_eg =fd.askdirectory()
    images_eg = image_read(path_eg)
    # cv2.imshow('1',images_eg[3])
    # cv2.waitKey(0)

    #获取分类图片名称
    image_names = []
    for i in range(len(images_eg)):
        image_name = i#这个部分要让用户输入名称
        image_names.append(image_name)
    # img_names = ["ljf1", "lsj"]
    #创建分类图片对应文件夹
    #path_rs = 'newimage'
    for i in range(len(images_eg)):
        os.makedirs(os.path.join(path_rs, str(image_names[i])))


    #读取待分类图片
    #path_cl = 'trainer'
    images_cl = image_read(path_cl)

    # 加载人脸检测器
    detector = dlib.get_frontal_face_detector()

    # 加载关键点检测器
    predictor = dlib.shape_predictor('D:/shape_predictor_68_face_landmarks.dat')

    # 加载人脸特征编码模型
    encoder = dlib.face_recognition_model_v1('D:/dlib_face_recognition_resnet_model_v1.dat')

    # 调用方法：128维特征向量输出
    imgs_eg_128D = []
    for i in range(len(images_eg)):
        img_128D = encoder_face(images_eg[i], detector, predictor, encoder)[0]
        imgs_eg_128D.append(img_128D)

    imgs_cl_128D = []
    for j in range(len(images_cl)):
        img_128D = encoder_face(images_cl[j], detector, predictor, encoder)[0]
        imgs_cl_128D.append(img_128D)


    # 调用方法：比较人脸，判断特征向量之间距离，判断是否为同一人,并存入对应文件夹
    imagePaths = [os.path.join(path_cl, f) for f in os.listdir(path_cl)]
    for i in range(len(images_cl)):
        distance, name = compare_faces_order(imgs_eg_128D, imgs_cl_128D[i], image_names)
        shutil.copy(imagePaths[i], os.path.join(path_rs, str(name[0])))
    messagebox.showinfo('提示', f'分类后的照片已保存到\n{path_cl}')

top = Tk()
top.title('dlib分类')  # 窗口标题
top.geometry('500x400')  # 窗口大小

show_folderPath1 = Entry(top)
show_folderPath1.grid(row=4, column=2,ipadx=30,padx=15,pady=10)
show_folderPath2 = Entry(top)
show_folderPath2.grid(row=2, column=2,ipadx=30,padx=15,pady=10)
show_folderPath3 = Entry(top)
show_folderPath3.grid(row=6, column=2,ipadx=30,padx=15,pady=10)
hdr3 = Label(top,text="请选择分类结果输出文件夹", font=("微软雅黑", 11)).grid(row=1, column=2,pady=5)
hdr2 = Label(top,text="请选择待分类人像照片所在的文件夹", font=("微软雅黑", 11)).grid(row=5, column=2,pady=5)
hdr1 = Label(top,text="请选择人像照片训练集文件夹", font=("微软雅黑", 11)).grid(row=3, column=2,pady=5)
btn1 = Button(top, text="浏览", command=openFloder1)  # 设置一个按钮
btn2 = Button(top, text="浏览", command=openFloder2)  # 设置一个按钮
btn3 = Button(top, text="浏览", command=openFloder3)  # 设置一个按钮
btn4 = Button(top, text="开始分类", command=main)  # 设置一个按钮
btn1.grid(row=4, column=3,ipadx=15,pady=10)
btn2.grid(row=2, column=3,ipadx=15,pady=10)
btn3.grid(row=6, column=3,ipadx=15,pady=10)
btn4.grid(row=7, column=2,ipadx=20,pady=20,ipady=5)
top.mainloop()




